#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep

import grpc

# Import P4Runtime lib from parent utils dir
# Probably there's a better way of doing this.
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

SWITCH_TO_HOST_PORT = 1

def get_switch_to_switch_port(switch_name, tunnel_id):
    """
    根据交换机名称和隧道ID确定交换机之间的连接端口
    """
    # s1 的端口配置
    if switch_name == 's1':
        if tunnel_id == 100:  # s1 -> s2 (h1 -> h2)
            return 2
        elif tunnel_id == 101:  # s1 -> s3 (h1 -> h3)
            return 3
    # s2 的端口配置        
    elif switch_name == 's2':
        if tunnel_id == 200:  # s2 -> s1 (h2 -> h1)
            return 2
        elif tunnel_id == 201:  # s2 -> s3 (h2 -> h3)
            return 3
    # s3 的端口配置
    elif switch_name == 's3':
        if tunnel_id == 300:  # s3 -> s1 (h3 -> h1)
            return 2
        elif tunnel_id == 301:  # s3 -> s2 (h3 -> h2)
            return 3
    else:
        raise ValueError(f"Unknown switch or tunnel ID: {switch_name}, {tunnel_id}")

def writeTunnelRules(p4info_helper, ingress_sw, egress_sw, tunnel_id,
                     dst_eth_addr, dst_ip_addr):
    """
    Installs three rules:
    1) An tunnel ingress rule on the ingress switch in the ipv4_lpm table that
       encapsulates traffic into a tunnel with the specified ID
    2) A transit rule on the ingress switch that forwards traffic based on
       the specified ID
    3) An tunnel egress rule on the egress switch that decapsulates traffic
       with the specified ID and sends it to the host

    :param p4info_helper: the P4Info helper
    :param ingress_sw: the ingress switch connection
    :param egress_sw: the egress switch connection
    :param tunnel_id: the specified tunnel ID
    :param dst_eth_addr: the destination IP to match in the ingress rule
    :param dst_ip_addr: the destination Ethernet address to write in the
                        egress rule
    """
    # 获取交换机间连接端口
    switch_to_switch_port = get_switch_to_switch_port(ingress_sw.name, tunnel_id)

    # 1) Tunnel Ingress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
        },
        action_name="MyIngress.myTunnel_ingress",
        action_params={
            "dst_id": tunnel_id,
        })
    ingress_sw.WriteTableEntry(table_entry)
    print(f"Installed ingress tunnel rule on {ingress_sw.name} for tunnel {tunnel_id} to {dst_ip_addr}")

    # 2) Tunnel Transit Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_forward",
        action_params={
            "port": switch_to_switch_port
        })
    ingress_sw.WriteTableEntry(table_entry)
    print(f"Installed transit tunnel rule on {ingress_sw.name} for tunnel {tunnel_id} to port {switch_to_switch_port}")

    # 3) Tunnel Egress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_egress",
        action_params={
            "dstAddr": dst_eth_addr,
            "port": SWITCH_TO_HOST_PORT
        })
    egress_sw.WriteTableEntry(table_entry)
    print(f"Installed egress tunnel rule on {egress_sw.name} for tunnel {tunnel_id} from {dst_eth_addr}")

def writeDirectRouteRules(p4info_helper, switch, dst_ip, dst_mac, port):
    """
    添加直接路由规则（备用方案）
    """
    try:
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.ipv4_lpm",
            match_fields={
                "hdr.ipv4.dstAddr": (dst_ip, 32)
            },
            action_name="MyIngress.ipv4_forward",
            action_params={
                "dstAddr": dst_mac,
                "port": port
            })
        switch.WriteTableEntry(table_entry)
        print(f"Installed direct route on {switch.name} for {dst_ip} to port {port}")
    except Exception as e:
        print(f"Failed to install direct route: {e}")

def readTableRules(p4info_helper, sw):
    """
    Reads the table entries from all tables on the switch.
    """
    print('\n----- Reading tables rules for %s -----' % sw.name)
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print('%s: ' % table_name, end=' ')
            for m in entry.match:
                print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                print('%r' % (p4info_helper.get_match_field_value(m),), end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print('->', action_name, end=' ')
            for p in action.params:
                print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                print('%r' % p.value, end=' ')
            print()

def printCounter(p4info_helper, sw, counter_name, index):
    """
    Reads the specified counter at the specified index from the switch.
    """
    for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
        for entity in response.entities:
            counter = entity.counter_entry
            print("%s %s %d: %d packets (%d bytes)" % (
                sw.name, counter_name, index,
                counter.data.packet_count, counter.data.byte_count
            ))

def printGrpcError(e):
    print("gRPC Error:", e.details(), end=' ')
    status_code = e.code()
    print("(%s)" % status_code.name, end=' ')
    traceback = sys.exc_info()[2]
    print("[%s:%d]" % (traceback.tb_frame.f_code.co_filename, traceback.tb_lineno))

def main(p4info_file_path, bmv2_file_path):
    # Instantiate a P4Runtime helper from the p4info file
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Create switch connections for s1, s2, and s3
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # Send master arbitration update message to establish this controller as master
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # Install the P4 program on all switches
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s1")
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s2")
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s3")

        print("\n=== Configuring Tunnel Rules ===")

        # 配置h1到h2的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s2, tunnel_id=100,
                         dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2")

        # 配置h1到h3的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s3, tunnel_id=101,
                         dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3")

        # 配置h2到h1的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s1, tunnel_id=200,
                         dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1")

        # 配置h2到h3的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s3, tunnel_id=201,
                         dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3")

        # 配置h3到h1的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s1, tunnel_id=300,
                         dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1")

        # 配置h3到h2的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s2, tunnel_id=301,
                         dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2")

        # 添加直接路由规则作为备用方案
        print("\n=== Configuring Direct Route Rules (Backup) ===")
        
        # s1 的直接路由
        writeDirectRouteRules(p4info_helper, s1, "10.0.2.2", "08:00:00:00:02:22", 2)  # h1->h2
        writeDirectRouteRules(p4info_helper, s1, "10.0.3.3", "08:00:00:00:03:33", 3)  # h1->h3
        
        # s2 的直接路由
        writeDirectRouteRules(p4info_helper, s2, "10.0.1.1", "08:00:00:00:01:11", 2)  # h2->h1
        writeDirectRouteRules(p4info_helper, s2, "10.0.3.3", "08:00:00:00:03:33", 3)  # h2->h3
        
        # s3 的直接路由
        writeDirectRouteRules(p4info_helper, s3, "10.0.1.1", "08:00:00:00:01:11", 2)  # h3->h1
        writeDirectRouteRules(p4info_helper, s3, "10.0.2.2", "08:00:00:00:02:22", 3)  # h3->h2

        # 读取所有交换机的表规则
        print("\n=== Reading Table Rules ===")
        readTableRules(p4info_helper, s1)
        readTableRules(p4info_helper, s2)
        readTableRules(p4info_helper, s3)

        print("\n=== Starting Counter Monitoring ===")
        # 打印隧道计数器
        while True:
            sleep(5)
            print('\n----- Reading tunnel counters -----')
            # h1相关的计数器
            print("=== h1 counters ===")
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 100)  # h1->h2
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 101)  # h1->h3
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 200)   # h2->h1
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 300)   # h3->h1
            
            # h2相关的计数器
            print("=== h2 counters ===")
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 200)  # h2->h1
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 201)  # h2->h3
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 100)   # h1->h2
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 301)   # h3->h2
            
            # h3相关的计数器
            print("=== h3 counters ===")
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 300)  # h3->h1
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 301)  # h3->h2
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 101)   # h1->h3
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 201)   # h2->h3

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
main(args.p4info, args.bmv2_json)